include "Math.ms"

--------------------------------------------------------------------------------
-- Structures
struct vBakingState
(
      mBakeeList
    , mCasterList
    , mLightList
    , mRMGIntersect
    , mCurrBakee
    , mStartTime
    
)

--------------------------------------------------------------------------------
-- Global Objects
--global gBakingState = undefined
global gBakeSamples = undefined

--------------------------------------------------------------------------------
-- Methods
function vBakerInit = 
(
    gBakeSamples = #()
    
    append gBakeSamples [ -0.94201624, -0.39906216 ]
    append gBakeSamples [ 0.94558609, -0.76890725 ]
    append gBakeSamples [ -0.094184101, -0.92938870 ]
    append gBakeSamples [ 0.34495938, 0.29387760 ]
    append gBakeSamples [ -0.91588581, 0.45771432 ]
    append gBakeSamples [ -0.81544232, -0.87912464 ]
    append gBakeSamples [ -0.38277543, 0.27676845 ]
    append gBakeSamples [ 0.97484398, 0.75648379 ]
    append gBakeSamples [ 0.44323325, -0.97511554 ]
    append gBakeSamples [ 0.53742981, -0.47373420 ]
    append gBakeSamples [ -0.26496911, -0.41893023 ]
    append gBakeSamples [ 0.79197514, 0.19090188 ]
    append gBakeSamples [ -0.24188840, 0.99706507 ]
    append gBakeSamples [ -0.81409955, 0.91437590 ]
    append gBakeSamples [ 0.19984126, 0.78641367 ]
    append gBakeSamples [ 0.14383161, -0.14100790 ]
)

function vBakerGetAmbient =
(
    local rgb = ( ambientColor as Point3 ) / 255.0
    return ( vGammaExpand rgb )
)

--------------------------------------------------------------------------------
function vBakerGetLightColor lt = 
(
    local rgb = ( lt.rgb as Point3 ) / 255.0
    return ( vGammaExpand rgb ) * lt.multiplier
)

--------------------------------------------------------------------------------
function vBakerGetLightVector lt pos =
(
    local lightDir
    
    if (   classof lt == OmniLight \
        or classof lt == FreeSpot \
        or classof lt == TargetSpot \
    ) then
    (
        lightDir = normalize ( lt.pos - pos )
    )
    else
    (
        -- default case
        lightDir = normalize ( lt.dir )
    )
    
    return lightDir
)

--------------------------------------------------------------------------------
function vBakerGenerateShadowSamples state r p n t b =
(
    local samples = #()
    
    local rgmi = state.mRMGIntersect
    
    local c = p + n * 0.1
    append samples ( c )
    
    for s in gBakeSamples do
    (
        local c2 = c + ( t * r * s.x ) + ( b * r * s.y )
        
        append samples ( c2 )
    )
    
    return samples
)

--------------------------------------------------------------------------------
function vBakerCalcSoftShadow rmgIntersect lt samples =
(
    local shadowTerm
    if lt.raytracedShadows == false then
    (
        shadowTerm = 1.0
    )
    else
    (
        if (   classof lt == OmniLight \
            or classof lt == FreeSpot \
            or classof lt == TargetSpot \
        ) then
        (
            -- point or spot light
            t = 0.0
            
            for s in samples do
            (
                if 0 == ( rmgIntersect.intersectSegment s lt.pos false ) do
                    t = t + 1.0
            )
            
            shadowTerm = t / samples.count
        )
        else
        (
            -- directional or other light types
            t = 0.0
            
            for s in samples do
            (
                if 0 == ( rmgIntersect.intersectRay s lt.dir false ) do
                    t = t + 1.0
            )
            
            shadowTerm = t / samples.count
        )
    )
    
    return shadowTerm
)

--------------------------------------------------------------------------------
function vBakerProcessSample state sampleRadius pos norm tangent binormal =
(
    local cnt = 1
    
    local lum = #()
    append lum ( vBakerGetAmbient() )
    
    local b = #()
    append b [0.0, 0.0, 1.0]
    
    local samples = ( vBakerGenerateShadowSamples state sampleRadius pos norm tangent binormal )
    
    for lt in state.mLightList do
    (
        --local shadowTerm = ( vBakerCalcSoftShadow state.mRMGIntersect lt samples )
        local shadowTerm = 1.0
        
        if shadowTerm > 0 do
        (
            local lightClr = vBakerGetLightColor lt
            local lightVec = vBakerGetLightVector lt pos
            
            lightVecTs = [0,0,0]
            lightVecTs.x = dot lightVec tangent
            lightVecTs.y = dot lightVec binormal
            lightVecTs.z = dot lightVec norm
            
            for i = 1 to cnt do
            (
                local bdotl = vSaturate( dot b[i] lightVecTs )
                lum[i] = lum[i] + lightClr * ( bdotl * shadowTerm )
            )
        )
    )
    
    --for i = 1 to cnt do
    --(
    --    lum[i] = ( vGammaCompress lum[i] )
    --)
    
    return lum
)

--------------------------------------------------------------------------------
function vBakerProcessSampleRNM state sampleRadius pos norm tangent binormal =
(
    local cnt = 3
    
    local lum = #()
    append lum ( vBakerGetAmbient() )
    append lum ( vBakerGetAmbient() )
    append lum ( vBakerGetAmbient() )
    
    local b = #()
    append b [-0.40824829046386307, 0.70710678118654746, 0.57735026918962584]
    append b [-0.40824829046386307, -0.70710678118654746, 0.57735026918962584]
    append b [0.81649658092772615, 0.0, 0.57735026918962584]
    
    local samples = ( vBakerGenerateShadowSamples state sampleRadius pos norm tangent binormal )
    
    for lt in state.mLightList do
    (
        local shadowTerm = ( vBakerCalcSoftShadow state.mRMGIntersect lt samples )
        
        if shadowTerm > 0 do
        (
            local lightClr = vBakerGetLightColor lt
            local lightVec = vBakerGetLightVector lt pos
            
            lightVecTs = [0,0,0]
            lightVecTs.x = dot lightVec tangent
            lightVecTs.y = dot lightVec binormal
            lightVecTs.z = dot lightVec norm
            
            for i = 1 to cnt do
            (
                local bdotl = vSaturate( dot b[i] lightVecTs )
                lum[i] = lum[i] + lightClr * ( bdotl * shadowTerm )
            )
        )
    )
    
    --for i = 1 to cnt do
    --(
    --    lum[i] = ( vGammaCompress lum[i] )
    --)
    
    return lum
)

--------------------------------------------------------------------------------
function vBakerProcessMesh state meshObj =
(
    format "processing % [ % vertices; % faces; % vcolour; % channels]...\n"\
        meshObj.name\
        ( getNumVerts meshObj )\
        ( getNumFaces meshObj )\
        ( getNumCPVVerts meshObj )\
        ( meshop.getNumMaps meshObj )
    
    local outChannel = 2
    local outCount = 1
    
    local outChannelBegin = outChannel
    local outChannelEnd = outChannel + outCount - 1
    
    if ( meshop.getNumMaps meshObj ) < ( outChannel + outCount ) do
    (
        format "%: output channels does not exists, building...\n" meshObj.name
        
        meshop.setNumMaps meshObj ( outChannel + outCount )
    )
    
    local faceCnt = ( getNumFaces meshObj )
    
    for channel = outChannelBegin to outChannelEnd do
    (
        format "reseting channel % of %...\n" channel meshObj.name
        
        meshop.setMapSupport meshObj channel true
        meshop.setNumMapFaces meshObj channel faceCnt
        meshop.setNumMapVerts meshObj channel ( faceCnt * 3 )
    )
    
    --local r = ( vGetNodeMaxSize meshObj ) * 0.05
    local r = 5.0
    
    local mfaceIdx = 1
    
    local gcCnt = 1
    
    -- write results back
    for fIdx = 1 to ( getNumFaces meshObj ) do
    (
        local face = ( getFace meshObj fIdx )
        
        local tface = meshop.getMapFace meshObj 1 fIdx
        
        local mface = [ mfaceIdx, mfaceIdx+1, mfaceIdx+2 ]
        
        for channel = outChannelBegin to outChannelEnd do
        (
            meshop.setMapFace meshObj channel fIdx mface
        )
        
        mfaceIdx = mfaceIdx + 3
        
        local p1 =  ( getVert meshObj face.x )
        local p2 =  ( getVert meshObj face.y )
        local p3 =  ( getVert meshObj face.z )
        
        local uv1 = ( meshop.getMapVert meshObj 1 tface.x )
        local uv2 = ( meshop.getMapVert meshObj 1 tface.y )
        local uv3 = ( meshop.getMapVert meshObj 1 tface.z )
        
        local tb = ( vCalcTangentAndBinormal p1 p2 p3 uv1 uv2 uv3 )
        
        local n = normalize( cross ( p2 - p1 ) ( p3 - p1 ) )
        
        for i = 1 to 3 do
        (
            local vIdx = face[i]
            
            local p = getVert meshObj vIdx
            --local n = getNormal meshObj vIdx
            local t = normalize tb[1]
            local b = normalize tb[2]
            
            t = normalize ( t - n * ( dot n t ) )
            b = normalize ( b - n * ( dot n b ) )
            
            --format "t: % b: % n: %\n" t b n
            
            --local results = ( vBakerProcessSampleRNM state r p n t b )
            local results = ( vBakerProcessSample state r p n t b )
            
            for channel = outChannelBegin to outChannelEnd do
            (
                local rIdx = channel - outChannelBegin + 1
                meshop.setMapVert meshObj channel mface[i] results[rIdx]
                --meshop.setMapVert meshObj channel mface[i] [0.0, 0.0, 0.0]
            )
            
            results = undefined
            
        )
        
        gcCnt = gcCnt + 1
        
        if ( gcCnt >= 16 ) do
        (
            gc()
            gcCnt = 1
        )
        --format "f: % vcf: %\n" face vcface
        --sleep( 1 )
    )
    
    gc()
    
    update meshObj
)

--------------------------------------------------------------------------------
function vBakerCreateState =
(
    state = vBakingState()
    
    state.mBakeeList = #()
    state.mCasterList = #()
    state.mLightList = #()
    state.mRMGIntersect = undefined
    
    return state
)

--------------------------------------------------------------------------------
function vBakerDestroyState state =
(
    state.mBakeeList = undefined
    state.mCasterList = undefined
    state.mRMGIntersect = undefined
    
    gc()
)

--------------------------------------------------------------------------------
function vBakerBeginBaking state gridSize =
(
    state.mCurrBakee = 1
    state.mStartTime = timeStamp()
    
    state.mRMGIntersect = RayMeshGridIntersect()
    state.mRMGIntersect.Initialize gridSize
    
    for caster in state.mCasterList do
    (
        state.mRMGIntersect.addNode caster
    )
    
    state.mRMGIntersect.buildGrid()
)

--------------------------------------------------------------------------------
function vBakerStepBaking state =
(
    if state.mCurrBakee <= state.mBakeeList.count do
    (
        local bakee = state.mBakeeList[ state.mCurrBakee ]
        
        vBakerProcessMesh state bakee
        
        state.mCurrBakee = state.mCurrBakee + 1
    )
    
    return state.mCurrBakee <= state.mBakeeList.count
)

--------------------------------------------------------------------------------
function vBakerEndBaking state =
(
    local endTime = timestamp()
    
    format "baking done in % seconds\n" ( ( endTime - state.mStartTime ) * 0.001 )
    
    state.mRMGIntersect = undefined
    state.mCurrBakee = undefined
    state.mStartTime = undefined
    
    gc()
)


--------------------------------------------------------------------------------
-- Test Function
function TestBaker =
(
    if selection.count == 0 then
    (
        print "Selection is empty!"
    )
    else
    (
        vBakerInit()
        
        local state = vBakerCreateState()
        
        for sel in selection do
        (
            if ( classof sel ) == Editable_mesh do
            (
                format "adding % to list\n" sel.name
                
                append state.mBakeeList( sel )
                append state.mCasterList( sel )
            )
        )
        
        if state.mBakeeList.count > 0 do
        (
            for obj in $* do
            (
                if isKindOf obj Light do
                (
                    format "adding % to light list\n" obj.name
                    append state.mLightList( obj )
                )
            )
            
            vBakerBeginBaking state 16
            
            while ( vBakerStepBaking state ) do
            (
                
            )
            
            vBakerEndBaking state
        )
        
        vBakerDestroyState state
        state =  undefined
    )
)

TestBaker()